import os

import torch

from P03_NER.LSTM_CRF.config import Config
from P03_NER.LSTM_CRF.model.BiLSTM import NERLSTM
from P03_NER.LSTM_CRF.model.BiLSTM_CRF import NERLSTM_CRF
from P03_NER.LSTM_CRF.utils.data_loader import word2id

conf = Config()
# 获取文件的绝对路径，然后基于这个路径进行路径拼接
base_dir = os.path.dirname(os.path.abspath(__file__))
print(f'base_dir-->{base_dir}')

# 1.实例化模型
models = {"BiLSTM": NERLSTM,
          "BiLSTM_CRF": NERLSTM_CRF}
model = models[conf.model](conf.embedding_dim, conf.hidden_dim, conf.dropout, word2id, conf.tag2id)
print(f'model--》{model}')
# 2.加载训练好的模型参数
if conf.model == "BiLSTM":
    model.load_state_dict(torch.load(os.path.join(base_dir, "save_model/bilstm_best.pth"), weights_only=True))
elif conf.model == "BiLSTM_CRF":
    model.load_state_dict(torch.load(os.path.join(base_dir, "save_model/bilstm_crf_best.pth"), weights_only=True))

# 通过tag2id构造id2tag
id2tag = {v: k for k, v in conf.tag2id.items()}
# print(f'id2tag--->{id2tag}')

def model2test(sample):
    # 3.处理数据
    # 1)文字转id
    ids = []
    for word in sample:
        if word not in word2id:
            word = 'UNK'
        ids.append(word2id[word])
    # print(f'ids--->{ids}')
    # 2)转成张量
    ids_tensor = torch.tensor([ids])  # 注意：需要添加batch_size维度
    # print(f'ids_tensor--->{ids_tensor}')

    # 3)生成attention_mask张量
    # attention_mask = torch.tensor([[1] * len(ids)])  # 方法一
    # print(f'attention_mask--->{attention_mask}')
    attention_mask = (ids_tensor !=0).long()  # 方法二
    # print(f'attention_mask--->{attention_mask}')

    # 4.模型预测
    model.eval()
    with torch.no_grad():
        if conf.model == "BiLSTM":
            # 送入模型
            pred = model(ids_tensor, attention_mask)
            # 通过argmax获取分数最大的索引
            predict = pred.argmax(dim=-1)[0]
            # print(f'predict--->{predict}')
            # 将标签id转成标签名称
            predict_tag = [id2tag[tag_id.item()] for tag_id in predict]
            # print(f'predict_tag--->{predict_tag}')

        elif conf.model == "BiLSTM_CRF":
            predict = model(ids_tensor, attention_mask)[0]
            # print(f'predict--->{predict}')
            # 将标签id转成标签名称
            predict_tag = [id2tag[tag_id] for tag_id in predict]
            # print(f'predict_tag--->{predict_tag}')

    # 5.结果处理
    # 将句子转成列表，方便后续处理
    chars = [word for word in sample]
    result = extract_entities(chars, predict_tag)
    return result


def extract_entities(text, tags):
    entities = {}
    entity = ""
    entity_type = ""

    for char, tag in zip(text, tags):
        if tag.startswith("B-"):
            # 遇到新实体，保存前一个实体
            if entity:
                entities[entity] = entity_type
            entity = char
            entity_type = tag[2:]  # 去掉 "B-"
        elif tag.startswith("I-") and entity:
            entity += char
        else:
            # 遇到 O 或不连续 I-，结束当前实体
            if entity:
                entities[entity] = entity_type
                entity = ""
                entity_type = ""

    # 结束时处理最后一个实体
    if entity:
        entities[entity] = entity_type

    return entities


if __name__ == '__main__':
    result = model2test(sample='入院后完善各项检查，给予右下肢持续皮牵引，应用健骨药物治疗，患者略发热，查血常规：白细胞数12.18*10^9/L，中性粒细胞百分比92.00%。给予应用抗生素预防感染。复查血常规：白细胞数6.55*10^9/L，中性粒细胞百分比74.70%，红细胞数2.92*10^12/L，血红蛋白94.0g/L。考虑贫血，指示加强营养。建议患者手术治疗,患者拒绝手术治疗。继续右下肢牵引，患者家属要求今日出院。')
    print(result)